
import json
from pathlib import Path
from secrets import token_bytes

def random_key(length):
    """根据指定长度生成随机密钥"""
    key = token_bytes(nbytes=length)
    key_int = int.from_bytes(key, 'big')
    return key_int

def encrypt(raw):
    """加密单元"""
    raw_bytes = raw.encode()
    raw_int = int.from_bytes(raw_bytes, 'big')
    key_int = random_key(len(raw_bytes))
    return raw_int ^ key_int, key_int

def decrypt(encrypted, key_int):
    """解密单元"""
    decrypted = encrypted ^ key_int
    length = (decrypted.bit_length() + 7) // 8
    decrypted_bytes = int.to_bytes(decrypted, length, 'big')
    return decrypted_bytes.decode()

def encrypt_file(path, key_path=None, *, encoding='utf-8'):
    """加密文件"""
    path = Path(path)
    cwd = path.parent / path.stem
    path_encrypted = cwd / (path.stem + '.enc')

    if key_path is None:
        key_path = cwd / 'key'

    if not cwd.exists():
        cwd.mkdir(parents=True, exist_ok=True)
        path_encrypted.touch()
        key_path.touch()

    # 打开文件并将内容写入
    with path.open('rt', encoding=encoding) as f1, \
         path_encrypted.open('wt', encoding=encoding) as f2, \
         key_path.open('wt', encoding=encoding) as f3:
        content = f1.read()
        encrypted, key = encrypt(content)
        json.dump(encrypted, f2)
        json.dump(key, f3)

def decrypt_file(path_encrypted, key_path=None, *, encoding='utf-8'):
    """解密文件"""
    path_encrypted = Path(path_encrypted)
    cwd = path_encrypted.parent
    path_decrypted = cwd / (path_encrypted.stem + '_decrypted')

    if not path_encrypted.exists():
        raise FileNotFoundError(f"Encrypted file '{path_encrypted}' does not exist.")

    if key_path is None:
        key_path = cwd / 'key'

    if not key_path.exists():
        raise FileNotFoundError(f"Key file '{key_path}' does not exist.")

    if not path_decrypted.exists():
        path_decrypted.mkdir(parents=True, exist_ok=True)
        path_decrypted /= path_encrypted.name
        path_decrypted.touch()

    # 打开文件并将内容写入
    with path_encrypted.open('rt', encoding=encoding) as f1, \
         key_path.open('rt', encoding=encoding) as f2, \
         path_decrypted.open('wt', encoding=encoding) as f3:
        encrypted = json.load(f1)
        key = json.load(f2)
        decrypted = decrypt(encrypted, key)
        f3.write(decrypted)

# 示例运行
if __name__ == '__main__':
    try:
        encrypt_file('/Users/liulu/PycharmProjects/practical_tools/util/input.txt')
        decrypt_file('/Users/liulu/PycharmProjects/practical_tools/util/input/input.enc')
    except FileNotFoundError as e:
        print(e)
    except Exception as e:
        print(f"An error occurred: {e}")